# coding: utf-8
# @File: dataset.py
# @Author: HE D.H.
# @Email: victor-he@qq.com
# @Time: 2021/12/09 11:01:32
# @Description:

import numpy as np
from torch.utils.data import Dataset
from tqdm import tqdm


class CNewsDataset(Dataset):
    def __init__(self, filename, tokenizer):
        # 数据集初始化
        self.labels = ['体育', '娱乐', '家居', '房产', '教育', '时尚', '时政', '游戏', '科技', '财经']
        self.labels_id = list(range(len(self.labels)))
        self.tokenizer = tokenizer
        self.input_ids = []
        self.token_type_ids = []
        self.attention_mask = []
        self.label_id = []
        self.load_data(filename)

    def load_data(self, filename):
        # 加载数据
        print('loading data from:', filename)
        with open(filename, 'r', encoding='utf-8') as rf:
            lines = rf.readlines()
        for line in tqdm(lines, ncols=100):
            label, text = line.strip().split('\t')
            label_id = self.labels.index(label)
            token = self.tokenizer(text,
                                   add_special_tokens=True,
                                   return_tensors="pt",  # 返回PyTorch张量
                                   padding='max_length',  # 将每个sequence填充到指定的最大长度。
                                   truncation=True,  # 如果为True，则每个序列中超过最大长度的标记将被截断。
                                   max_length=512  # 每个sequence的最大长度
                                   )
            '''
                > input_ids：文本对应的 token ID 序列
                > token_type_ids:区分句子（0 代表第一个句子，1 代表第二个句子）
                > 标记有效 token（1）和填充 token（0）
            '''
            self.input_ids.append(np.array(token['input_ids']).squeeze(0))
            self.token_type_ids.append(np.array(token['token_type_ids']).squeeze(0))
            self.attention_mask.append(np.array(token['attention_mask']).squeeze(0))
            self.label_id.append(label_id)

    def __getitem__(self, index):
        return self.input_ids[index], self.token_type_ids[index], self.attention_mask[index], self.label_id[index]

    def __len__(self):
        return len(self.input_ids)
